import argparse
from tqdm import tqdm
from typing import List
from pathlib import Path
from src.huffman_extractor import HuffmanExtractor as Extractor
from src.huffman_compressor import HuffmanCompressor as Compressor

class CompressionTool:
    def __init__(self, input_files: List[Path], output_dir: Path, action: str, overwrite: bool):
        self.input_files = input_files
        self.output_dir = output_dir
        self.action = action
        self.overwrite = overwrite

    def _compress(self, file: Path, output_path: Path) -> None:
        """文件压缩"""
        print(f"压缩文件：{file} -> {output_path}")
        compressor = Compressor(file, output_path)
        compressor.compress()

    def _extract(self, file: Path, output_path: Path) -> None:
        """文件解压"""
        print(f"解压文件：{file} -> {output_path}")
        extractor = Extractor(file, output_path)
        extractor.extract()

    def process_files(self) -> None:
        """处理输入文件，根据模式进行压缩或解压"""
        if not self.output_dir.exists():
            print(f"创建输出目录：{self.output_dir}")
            self.output_dir.mkdir(parents=True)

        # 使用 tqdm 创建进度条
        with tqdm(total=len(self.input_files), desc="处理文件", unit="file") as pbar:
            for file in self.input_files:
                if not file.exists():
                    print(f"文件不存在，跳过：{file}")
                    pbar.update(1)
                    continue

                if self.action == 'compress':
                    output_file = self.output_dir / f"compressed_{file.name}"
                    if output_file.exists() and not self.overwrite:
                        print(f"文件已存在，跳过：{output_file}")
                        pbar.update(1)
                        continue
                    self._compress(file, output_file)

                elif self.action == 'extract':
                    output_file = self.output_dir / f"extracted_{file.name}"
                    if output_file.exists() and not self.overwrite:
                        print(f"文件已存在，跳过：{output_file}")
                        pbar.update(1)
                        continue
                    self._extract(file, output_file)

                pbar.update(1)
    
    def interact(self) -> None:
        """交互模式"""
        while True:
            print("\n请输入一个目录路径（或者输入exit退出）：")
            dir_path = input().strip()
            if dir_path.lower() == 'exit':
                print("退出交互模式")
                break

            dir_path = Path(dir_path)
            if not dir_path.is_dir():
                print(f"{dir_path} 不是有效的目录，请重新输入！")
                continue

            txt_files = list(dir_path.glob("*.txt"))
            if not txt_files:
                print("该目录下没有 .txt 文件。")
                continue

            print("\n可操作的文件：")
            for idx, file in enumerate(txt_files, start=1):
                print(f"{idx}. {file.name}")

            print("\n请选择文件，输入文件的序号（多个序号用逗号分隔）：")
            selected_files_idx = input().strip()
            selected_files_idx = selected_files_idx.split(',')

            selected_files = []
            for idx in selected_files_idx:
                try:
                    selected_files.append(txt_files[int(idx.strip()) - 1])
                except (ValueError, IndexError):
                    print(f"无效的序号：{idx}")

            if not selected_files:
                print("没有选择有效的文件，重新开始。")
                continue

            for file in selected_files:
                print(f"\n请选择针对 {file.name} 的操作：1. 压缩  2. 解压")
                action_choice = input().strip()
                action = 'compress' if action_choice == '1' else 'extract'

                print("\n请输入保存路径（包括文件名）：")
                save_path = input().strip()
                save_path = Path(save_path)

                if action == 'compress':
                    self._compress(file, save_path)
                elif action == 'extract':
                    self._extract(file, save_path)

def main() -> None:
    parser = argparse.ArgumentParser(
        description="逐个文件压缩/解压工具",
        epilog="示例：my_tool -c -i file1.txt file2.txt -o ./output -y"
    )

    # 操作模式
    action_group = parser.add_mutually_exclusive_group()
    action_group.add_argument('-c', '--compress', action='store_true', help="压缩模式")
    action_group.add_argument('-x', '--extract', action='store_true', help="解压模式")

    # 输入文件
    parser.add_argument(
        '-i', '--input', 
        nargs='+', 
        type=Path, 
        help="输入文件，支持多个文件"
    )

    # 输出路径
    parser.add_argument(
        '-o', '--output',
        type=Path,
        default=Path.cwd(),
        help="输出路径（目录），默认为当前目录"
    )

    # 是否覆盖
    overwrite_group = parser.add_mutually_exclusive_group()
    overwrite_group.add_argument(
        '-y', '--yes', 
        action='store_true', 
        help="覆盖已存在文件"
    )
    overwrite_group.add_argument(
        '-n', '--no', 
        action='store_true', 
        help="禁止覆盖已存在文件"
    )

    # 交互模式
    parser.add_argument(
        '--interact', 
        action='store_true', 
        help="启用交互模式"
    )

    args = parser.parse_args()

    # 选择覆盖逻辑
    overwrite = args.yes or not args.no
    action = 'compress' if args.compress else 'extract'

    if args.interact:
        tool = CompressionTool(
            input_files=[],
            output_dir=Path.cwd(),
            action='compress', 
            overwrite=overwrite
        )
        tool.interact()
    else:
        tool = CompressionTool(
            input_files=args.input,
            output_dir=args.output,
            action=action,
            overwrite=overwrite
        )
        tool.process_files()

if __name__ == "__main__":
    main()